#!/usr/bin/env python3

import os
import sys
import glob
import re
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.patches import Rectangle
from pathlib import Path
from typing import Dict, List, Any

# --- 1. 配置路径 ---
# (无变化)
BASE_DIR = Path("/Users/chenxi/Documents/比赛/CCF/checker/example")
IN_DIR = BASE_DIR / "infile"
OUT_DIR = BASE_DIR / "outfile"
FIG_DIR = BASE_DIR / "fig"

# --- 2. 辅助函数：解析文件 ---
# (无变化)

def parse_infile(filepath: Path) -> Dict[str, Any]:
    """解析 infile，返回 L, M, N 和一个请求字典。"""
    requests = {}
    with open(filepath, 'r') as f:
        line = f.readline().split()
        L, M, N = int(line[0]), int(line[1]), int(line[2])
        
        for i in range(N):
            line = f.readline().split()
            requests[i] = {
                'addr': int(line[0]),
                'size': int(line[1]),
                'start': int(line[2]),
                'time': int(line[3]),
                'id': i
            }
    return {'L': L, 'M': M, 'N': N, 'requests': requests}

def parse_outfile(filepath: Path) -> List[Dict[str, Any]]:
    """解析 outfile，返回一个操作列表。"""
    ops = []
    with open(filepath, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if not parts:
                continue
            
            op_name = parts[0]
            T = int(parts[1])
            op = {'op': op_name, 'T': T}
            
            if op_name in ("Reload", "Offload"):
                op['A'] = int(parts[2])
                op['S'] = int(parts[3])
            elif op_name == "Visit":
                op['A'] = int(parts[2]) # 这是任务 ID
            
            ops.append(op)
    return ops

# --- 3. [!!! 核心修改 !!!] 甘特图可视化函数 ---

def visualize_case(infile_path: Path, outfile_path: Path, fig_path: Path):
    """
    为单个用例生成并保存 甘特图 可视化。
    [!!! MODIFIED !!!]
    图像包含三个子图：
    1. "Requests" (来自 Infile 的理论时间)
    2. "Visits" (NPU 实际调度时间)
    3. "IO" (Reload/Offload 串行时间)
    """
    
    try:
        in_data = parse_infile(infile_path)
        out_data = parse_outfile(outfile_path)
    except Exception as e:
        print(f"  [!] 失败: 解析文件时出错 {e}")
        return

    L, M, N = in_data['L'], in_data['M'], in_data['N']
    requests = in_data['requests']
    
    # 模拟 HBM 状态 (不变)
    try:
        hbm_state = np.zeros(L, dtype=np.int8)
    except MemoryError:
        print(f"  [!] 失败: 无法为 L={L} 分配内存。跳过...")
        return
        
    # [!!! NEW !!!] 存储 "Requests" (理论) 时间窗口
    request_bars = []
    
    visit_bars = [] # 存储 (task_id, T_start, duration, label, color)
    io_bars = []    # 存储 (T_start, duration, label, color)
    T_fin = 0
    
    colors = list(mcolors.TABLEAU_COLORS.values())
    
    # [!!! NEW !!!]
    # 遍历 Infile requests 来收集 "Requests" 绘图数据
    for task_id, req in requests.items():
        color = colors[task_id % len(colors)]
        request_bars.append((
            task_id,
            req['start'],  # 理论最早开始时间
            req['time'],   # 理论执行时间
            f"R{task_id}",  # 标签 (Request)
            color
        ))
        # 确保 T_fin 至少和理论时间一样长
        T_fin = max(T_fin, req['start'] + req['time'])


    # 遍历 Outfile (不变)
    for op in out_data:
        T_start = op['T']
        T_fin = max(T_fin, T_start)
        
        if op['op'] == 'Reload':
            A, S = op['A'], op['S']
            try:
                chunk = hbm_state[A : A+S]
                new_bytes = (chunk == 0)
                actual_reload_size = int(np.sum(new_bytes))
                chunk[new_bytes] = 1 
            except ValueError:
                continue 

            duration = 40 * actual_reload_size
            if duration > 0:
                io_bars.append((T_start, duration, 'R', colors[0])) # Blue
                T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Offload':
            A, S = op['A'], op['S']
            try:
                chunk = hbm_state[A : A+S]
                offload_bytes = (chunk == 1)
                actual_offload_size = int(np.sum(offload_bytes))
                chunk[offload_bytes] = 0
            except ValueError:
                continue

            duration = 40 * actual_offload_size
            if duration > 0:
                io_bars.append((T_start, duration, 'O', colors[3])) # Red
                T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Visit':
            task_id = op['A']
            if task_id not in requests:
                continue
                
            req = requests[task_id]
            duration = req['time']
            color = colors[task_id % len(colors)]
            
            visit_bars.append((task_id, T_start, duration, f"V{task_id}", color))
            T_fin = max(T_fin, T_start + duration)

        elif op['op'] == 'Fin':
            T_fin = max(T_fin, T_start)
            break
            
    # --- 4. 绘图 ---
    
    # [!!! MODIFIED !!!]
    # 高度比例：为 N 个 Requests, N 个 Visits 和 1 个 IO 任务分配空间
    height_ratio = [max(1, N), max(1, N), 1.5] # Requests, Visits, IO
    
    fig, (ax_requests, ax_visits, ax_io) = plt.subplots(
        3, 1, # [!!! MODIFIED !!!] 3 个子图
        figsize=(20, 3 + N * 1.0 + 1.0), # [!!! MODIFIED !!!] 动态计算总高度
        sharex=True, 
        gridspec_kw={'height_ratios': height_ratio}
    )
    
    # [!!! MODIFIED !!!] 移除图名
    # fig.suptitle(f"Scheduling Gantt Chart: {outfile_path.name}", fontsize=16)

    # [!!! NEW !!!]
    # --- 子图 1: "Requests" (Infile 理论时间) ---
    ax_requests.set_ylabel("Requests (Input)")
    
    for (task_id, T_start, duration, label, color) in request_bars:
        ax_requests.barh(
            y=task_id, 
            width=duration, 
            left=T_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.4,  # 使用较低的 alpha
            color=color,
            hatch='//'  # 使用斜线填充以区分
        )
        # 在条形图内部添加标签
        ax_requests.text(
            T_start + duration/2, 
            task_id, 
            label, 
            ha='center', 
            va='center', 
            color='black', 
            alpha=0.7,
            fontsize=9
        )
        
    # Y 轴设置 (让 Lane 0 在顶部)
    ax_requests.set_yticks(range(N))
    ax_requests.set_yticklabels([f"Lane {i}" for i in range(N)])
    ax_requests.set_ylim(N - 0.5, -0.5)
    ax_requests.grid(axis='x', linestyle=':', alpha=0.6)

    # --- 子图 2: "Visits" (NPU 实际调度) ---
    # [!!! MODIFIED !!!]
    ax_visits.set_ylabel("Visits (Scheduled)")
    
    for (task_id, T_start, duration, label, color) in visit_bars:
        ax_visits.barh(
            y=task_id, 
            width=duration, 
            left=T_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.8, # [!!! MODIFIED !!!] 较高的 alpha
            color=color
        )
        # 在条形图内部添加标签
        ax_visits.text(
            T_start + duration/2, 
            task_id, 
            label, 
            ha='center', 
            va='center', 
            color='white', 
            weight='bold',
            fontsize=10
        )
        
    # Y 轴设置 (让 Lane 0 在顶部)
    ax_visits.set_yticks(range(N))
    ax_visits.set_yticklabels([f"Lane {i}" for i in range(N)])
    ax_visits.set_ylim(N - 0.5, -0.5)
    ax_visits.grid(axis='x', linestyle=':', alpha=0.6)

    # --- 子图 3: "IO" (Reload/Offload) ---
    # (无变化)
    ax_io.set_ylabel("IO (Serial)")
    
    for (T_start, duration, label, color) in io_bars:
        ax_io.barh(
            y=0, 
            width=duration, 
            left=T_start, 
            height=0.8, 
            edgecolor='black', 
            alpha=0.75,
            color=color
        )
        # 标签
        ax_io.text(
            T_start + duration/2, 
            0, 
            label, 
            ha='center', 
            va='center', 
            color='white', 
            weight='bold',
            fontsize=10
        )

    # Y 轴设置
    ax_io.set_yticks([0])
    ax_io.set_yticklabels(["IO Lane 0"])
    ax_io.set_ylim(-0.5, 0.5)
    ax_io.grid(axis='x', linestyle=':', alpha=0.6)

    # --- 5. 保存 ---
    ax_io.set_xlabel("Time (T)")
    ax_io.set_xlim(left=-T_fin * 0.01, right=T_fin * 1.05)
    
    # [!!! MODIFIED !!!] 调整 tight_layout 以适应无标题
    plt.tight_layout(rect=[0, 0.03, 1, 1.0]) 
    plt.savefig(fig_path)
    plt.close(fig)
    print(f"  [✓] 成功: 已保存到 {fig_path.name}")


# --- 6. 主程序 ---
# (无变化)
def main():
    print(f"--- 开始可视化 (Gantt 模式) ---")
    print(f"输入目录: {IN_DIR}")
    print(f"输出目录: {OUT_DIR}")
    print(f"图形目录: {FIG_DIR}")
    print("--------------------")

    FIG_DIR.mkdir(parents=True, exist_ok=True)
    
    infile_paths = sorted(glob.glob(str(IN_DIR / "infile_*.txt")))
    
    if not infile_paths:
        print(f"[!] 错误: 在 {IN_DIR} 中未找到 'infile_*.txt' 文件。")
        return

    for infile_path in infile_paths:
        infile_path = Path(infile_path)
        
        match = re.search(r'infile_(\w+)\.txt', infile_path.name)
        if not match:
            print(f"[!] 跳过: 无法解析文件名 {infile_path.name}")
            continue
            
        file_id = match.group(1)
        outfile_name = f"outfile_{file_id}.txt"
        outfile_path = OUT_DIR / outfile_name
        
        fig_name = f"fig_{file_id}_gantt.png" 
        fig_path = FIG_DIR / fig_name
        
        print(f"[*] 正在处理: {infile_path.name} -> {outfile_name}")

        if not outfile_path.exists():
            print(f"  [!] 跳过: 未找到对应的 {outfile_name}")
            continue
            
        try:
            visualize_case(infile_path, outfile_path, fig_path)
        except Exception as e:
            print(f"  [!] 失败: 处理 {infile_path.name} 时发生意外错误: {e}")

    print("--------------------")
    print("--- 可视化完成 ---")

if __name__ == "__main__":
    main()